"""
Cython增强比对器

使用Cython优化的高性能比对器，完全替代现有的NumPy版本
"""

import logging
import time
import gc
import numpy as np
import pandas as pd
from typing import List, Dict, Tuple, Optional, Any, Callable
from concurrent.futures import ProcessPoolExecutor, as_completed
import multiprocessing as mp

from cython_integration import get_cython_engine, CythonCompareEngine


class CythonEnhancedComparer:
    """Cython增强比对器
    
    使用Cython优化的核心算法，提供最高性能的帧比对功能
    """
    
    def __init__(self, max_diff_reports: int = 100, use_parallel: bool = True):
        """初始化Cython增强比对器
        
        Args:
            max_diff_reports: 最大差异报告数量
            use_parallel: 是否使用并行处理
        """
        self.logger = logging.getLogger('cython_enhanced_comparer')
        self.max_diff_reports = max_diff_reports
        self.use_parallel = use_parallel
        
        # 获取Cython引擎
        self.cython_engine = get_cython_engine(max_diff_reports, use_parallel)
        
        # 设置处理策略
        self.large_file_threshold = 10000  # 大文件阈值
        self.batch_size = 1000  # 批处理大小
        
        self.logger.info(f"Cython增强比对器初始化完成 - Cython可用: {self.cython_engine.is_available()}")
    
    def compare_files_enhanced(self, 
                             frames1: List[dict], 
                             frames2: List[dict],
                             compare_range: Optional[Tuple[int, int]] = None,
                             progress_callback: Optional[Callable] = None,
                             file1_format: str = '1092',
                             file2_format: str = '1092') -> Dict:
        """增强的文件比对功能
        
        Args:
            frames1: 基准文件的帧列表
            frames2: 目标文件的帧列表
            compare_range: 比对范围
            progress_callback: 进度回调函数
            file1_format: 基准文件格式
            file2_format: 目标文件格式
            
        Returns:
            dict: 比对结果
        """
        start_time = time.time()
        
        try:
            self.logger.info(f"开始Cython增强比对: {len(frames1)} vs {len(frames2)} 帧")
            
            # 初始化结果结构
            comparison_result = {
                'total_frames': [len(frames1), len(frames2)],
                'frame_results': [],
                'matching_frames': 0,
                'virtual_channel_results': {},
                'processing_time': 0,
                'performance_info': self.cython_engine.get_performance_info(),
                'file1_format': file1_format,
                'file2_format': file2_format
            }
            
            if progress_callback:
                progress_callback(10, "准备数据...", "提取虚拟信道信息")
            
            # 快速提取虚拟信道信息
            start_extract = time.time()
            vc_ids1, frame_counts1 = self.cython_engine.extract_virtual_channel_info_fast(frames1, file1_format)
            vc_ids2, frame_counts2 = self.cython_engine.extract_virtual_channel_info_fast(frames2, file2_format)
            extract_time = time.time() - start_extract
            
            self.logger.debug(f"虚拟信道信息提取完成: {extract_time:.3f}秒")
            
            if progress_callback:
                progress_callback(25, "分组数据...", "按虚拟信道分组")
            
            # 按虚拟信道分组
            vc_groups1, vc_groups2 = self._group_frames_by_vc_fast(
                frames1, frames2, vc_ids1, vc_ids2, frame_counts1, frame_counts2
            )
            
            if progress_callback:
                progress_callback(40, "开始比对...", f"发现{len(vc_groups1)}个虚拟信道")
            
            # 选择比对策略
            total_comparisons = sum(
                min(len(vc_groups1.get(vc_id, [])), len(vc_groups2.get(vc_id, []))) 
                for vc_id in set(vc_groups1.keys()) | set(vc_groups2.keys())
            )
            
            if total_comparisons > self.large_file_threshold and self.use_parallel:
                # 大文件并行处理
                vc_results = self._parallel_compare_virtual_channels(
                    vc_groups1, vc_groups2, compare_range, progress_callback
                )
            else:
                # 串行处理
                vc_results = self._sequential_compare_virtual_channels(
                    vc_groups1, vc_groups2, compare_range, progress_callback
                )
            
            # 合并结果
            comparison_result['virtual_channel_results'] = vc_results
            
            # 统计信息
            total_matching = sum(vc_result.get('matching_frames', 0) for vc_result in vc_results.values())
            comparison_result['matching_frames'] = total_matching
            
            # 生成帧结果列表（用于兼容性）
            frame_results = []
            for vc_result in vc_results.values():
                frame_results.extend(vc_result.get('frame_comparisons', []))
            comparison_result['frame_results'] = frame_results
            
            elapsed_time = time.time() - start_time
            comparison_result['processing_time'] = elapsed_time
            
            if progress_callback:
                progress_callback(100, "比对完成", f"耗时{elapsed_time:.2f}秒")
            
            self.logger.info(f"Cython增强比对完成: 匹配{total_matching}帧, 耗时{elapsed_time:.2f}秒")
            
            return comparison_result
            
        except Exception as e:
            self.logger.error(f"Cython增强比对失败: {str(e)}", exc_info=True)
            raise
    
    def _group_frames_by_vc_fast(self, frames1, frames2, vc_ids1, vc_ids2, 
                                frame_counts1, frame_counts2) -> Tuple[Dict, Dict]:
        """快速按虚拟信道分组帧数据"""
        vc_groups1 = {}
        vc_groups2 = {}
        
        # 分组基准文件
        for i, (frame, vc_id, frame_count) in enumerate(zip(frames1, vc_ids1, frame_counts1)):
            if vc_id not in vc_groups1:
                vc_groups1[vc_id] = []
            
            frame_with_vc = frame.copy()
            frame_with_vc['virtual_channel'] = vc_id
            frame_with_vc['vc_frame_count'] = frame_count
            frame_with_vc['_index'] = i
            
            vc_groups1[vc_id].append(frame_with_vc)
        
        # 分组目标文件
        for i, (frame, vc_id, frame_count) in enumerate(zip(frames2, vc_ids2, frame_counts2)):
            if vc_id not in vc_groups2:
                vc_groups2[vc_id] = []
            
            frame_with_vc = frame.copy()
            frame_with_vc['virtual_channel'] = vc_id
            frame_with_vc['vc_frame_count'] = frame_count
            frame_with_vc['_index'] = i
            
            vc_groups2[vc_id].append(frame_with_vc)
        
        return vc_groups1, vc_groups2
    
    def _sequential_compare_virtual_channels(self, vc_groups1, vc_groups2, 
                                           compare_range, progress_callback) -> Dict:
        """串行比对虚拟信道"""
        vc_results = {}
        all_vc_ids = sorted(set(vc_groups1.keys()) | set(vc_groups2.keys()))
        
        for i, vc_id in enumerate(all_vc_ids):
            frames1_in_vc = vc_groups1.get(vc_id, [])
            frames2_in_vc = vc_groups2.get(vc_id, [])
            
            if progress_callback and i % 5 == 0:
                progress = 40 + int((i / len(all_vc_ids)) * 50)
                progress_callback(progress, f"比对虚拟信道{vc_id}", f"进度{i}/{len(all_vc_ids)}")
            
            vc_result = self._compare_single_vc_enhanced(
                vc_id, frames1_in_vc, frames2_in_vc, compare_range
            )
            vc_results[vc_id] = vc_result
        
        return vc_results
    
    def _parallel_compare_virtual_channels(self, vc_groups1, vc_groups2, 
                                         compare_range, progress_callback) -> Dict:
        """并行比对虚拟信道"""
        vc_results = {}
        all_vc_ids = sorted(set(vc_groups1.keys()) | set(vc_groups2.keys()))
        
        # 计算最优进程数
        num_processes = min(mp.cpu_count(), len(all_vc_ids), 8)
        
        if num_processes <= 1:
            return self._sequential_compare_virtual_channels(vc_groups1, vc_groups2, compare_range, progress_callback)
        
        self.logger.info(f"使用{num_processes}个进程并行比对{len(all_vc_ids)}个虚拟信道")
        
        try:
            with ProcessPoolExecutor(max_workers=num_processes) as executor:
                # 提交任务
                future_to_vc = {}
                for vc_id in all_vc_ids:
                    frames1_in_vc = vc_groups1.get(vc_id, [])
                    frames2_in_vc = vc_groups2.get(vc_id, [])
                    
                    future = executor.submit(
                        _compare_vc_worker, 
                        vc_id, frames1_in_vc, frames2_in_vc, 
                        compare_range, self.max_diff_reports
                    )
                    future_to_vc[future] = vc_id
                
                # 收集结果
                completed = 0
                for future in as_completed(future_to_vc):
                    vc_id = future_to_vc[future]
                    try:
                        vc_result = future.result()
                        vc_results[vc_id] = vc_result
                        completed += 1
                        
                        if progress_callback:
                            progress = 40 + int((completed / len(all_vc_ids)) * 50)
                            progress_callback(progress, f"完成虚拟信道{vc_id}", f"进度{completed}/{len(all_vc_ids)}")
                            
                    except Exception as e:
                        self.logger.error(f"虚拟信道{vc_id}比对失败: {e}")
                        vc_results[vc_id] = {
                            'vc_id': vc_id,
                            'error': str(e),
                            'frames_in_base': len(vc_groups1.get(vc_id, [])),
                            'frames_in_target': len(vc_groups2.get(vc_id, []))
                        }
        
        except Exception as e:
            self.logger.error(f"并行比对失败，回退到串行: {e}")
            return self._sequential_compare_virtual_channels(vc_groups1, vc_groups2, compare_range, progress_callback)
        
        return vc_results
    
    def _compare_single_vc_enhanced(self, vc_id: int, frames1: List[Dict], 
                                   frames2: List[Dict], compare_range: Optional[Tuple[int, int]]) -> Dict:
        """使用Cython增强算法比对单个虚拟信道"""
        vc_result = {
            'vc_id': vc_id,
            'frames_in_base': len(frames1),
            'frames_in_target': len(frames2),
            'matching_frames': 0,
            'mismatching_frames': 0,
            'missing_in_target': 0,
            'missing_in_base': 0,
            'frame_comparisons': []
        }
        
        if not frames1 and not frames2:
            return vc_result
        
        # 创建帧计数映射
        frames1_dict = {frame['vc_frame_count']: frame for frame in frames1}
        frames2_dict = {frame['vc_frame_count']: frame for frame in frames2}
        
        # 获取所有帧计数
        all_frame_counts = sorted(set(frames1_dict.keys()) | set(frames2_dict.keys()))
        
        # 批量比对策略
        if len(all_frame_counts) > self.batch_size:
            return self._batch_compare_vc_frames(
                vc_id, frames1_dict, frames2_dict, all_frame_counts, compare_range
            )
        else:
            return self._simple_compare_vc_frames(
                vc_id, frames1_dict, frames2_dict, all_frame_counts, compare_range
            )
    
    def _batch_compare_vc_frames(self, vc_id: int, frames1_dict: Dict, frames2_dict: Dict,
                                all_frame_counts: List[int], compare_range: Optional[Tuple[int, int]]) -> Dict:
        """批量比对虚拟信道帧"""
        vc_result = {
            'vc_id': vc_id,
            'frames_in_base': len(frames1_dict),
            'frames_in_target': len(frames2_dict),
            'matching_frames': 0,
            'mismatching_frames': 0,
            'missing_in_target': 0,
            'missing_in_base': 0,
            'frame_comparisons': []
        }
        
        # 分批处理
        for i in range(0, len(all_frame_counts), self.batch_size):
            batch_counts = all_frame_counts[i:i + self.batch_size]
            
            # 准备批量数据
            batch_payloads1 = []
            batch_payloads2 = []
            batch_frame_info = []
            
            for frame_count in batch_counts:
                frame1 = frames1_dict.get(frame_count)
                frame2 = frames2_dict.get(frame_count)
                
                if frame1 is None:
                    vc_result['missing_in_base'] += 1
                    vc_result['frame_comparisons'].append({
                        'frame_count': frame_count,
                        'type': 'missing_in_base',
                        'vc_id': vc_id
                    })
                    continue
                    
                if frame2 is None:
                    vc_result['missing_in_target'] += 1
                    vc_result['frame_comparisons'].append({
                        'frame_count': frame_count,
                        'type': 'missing_in_target',
                        'vc_id': vc_id
                    })
                    continue
                
                # 提取比对数据
                data1 = self.cython_engine.extract_comparison_data(frame1, compare_range)
                data2 = self.cython_engine.extract_comparison_data(frame2, compare_range)
                
                batch_payloads1.append(data1)
                batch_payloads2.append(data2)
                batch_frame_info.append((frame_count, frame1, frame2))
            
            # 批量比对
            if batch_payloads1:
                batch_results = self.cython_engine.batch_compare_payloads(
                    batch_payloads1, batch_payloads2, compare_range
                )
                
                # 处理批量结果
                for j, (is_match, (frame_count, frame1, frame2)) in enumerate(zip(batch_results, batch_frame_info)):
                    if is_match:
                        vc_result['matching_frames'] += 1
                        vc_result['frame_comparisons'].append({
                            'frame_count': frame_count,
                            'type': 'matching',
                            'vc_id': vc_id
                        })
                    else:
                        vc_result['mismatching_frames'] += 1
                        
                        # 详细比对不匹配的帧
                        detailed_result = self.cython_engine.compare_frames_detailed(
                            frame1, frame2, compare_range
                        )
                        
                        frame_comp = {
                            'frame_count': frame_count,
                            'type': 'content_mismatch',
                            'vc_id': vc_id
                        }
                        
                        if 'differences' in detailed_result:
                            frame_comp['differences'] = detailed_result['differences']
                        
                        vc_result['frame_comparisons'].append(frame_comp)
        
        return vc_result
    
    def _simple_compare_vc_frames(self, vc_id: int, frames1_dict: Dict, frames2_dict: Dict,
                                 all_frame_counts: List[int], compare_range: Optional[Tuple[int, int]]) -> Dict:
        """简单比对虚拟信道帧（小批量）"""
        vc_result = {
            'vc_id': vc_id,
            'frames_in_base': len(frames1_dict),
            'frames_in_target': len(frames2_dict),
            'matching_frames': 0,
            'mismatching_frames': 0,
            'missing_in_target': 0,
            'missing_in_base': 0,
            'frame_comparisons': []
        }
        
        for frame_count in all_frame_counts:
            frame1 = frames1_dict.get(frame_count)
            frame2 = frames2_dict.get(frame_count)
            
            if frame1 is None:
                vc_result['missing_in_base'] += 1
                vc_result['frame_comparisons'].append({
                    'frame_count': frame_count,
                    'type': 'missing_in_base',
                    'vc_id': vc_id
                })
                continue
                
            if frame2 is None:
                vc_result['missing_in_target'] += 1
                vc_result['frame_comparisons'].append({
                    'frame_count': frame_count,
                    'type': 'missing_in_target',
                    'vc_id': vc_id
                })
                continue
            
            # 详细比对
            detailed_result = self.cython_engine.compare_frames_detailed(
                frame1, frame2, compare_range
            )
            
            if detailed_result.get('type') == 'matching':
                vc_result['matching_frames'] += 1
                vc_result['frame_comparisons'].append({
                    'frame_count': frame_count,
                    'type': 'matching',
                    'vc_id': vc_id
                })
            else:
                vc_result['mismatching_frames'] += 1
                frame_comp = {
                    'frame_count': frame_count,
                    'type': 'content_mismatch',
                    'vc_id': vc_id
                }
                
                if 'differences' in detailed_result:
                    frame_comp['differences'] = detailed_result['differences']
                
                vc_result['frame_comparisons'].append(frame_comp)
        
        return vc_result


def _compare_vc_worker(vc_id: int, frames1: List[Dict], frames2: List[Dict], 
                      compare_range: Optional[Tuple[int, int]], max_diff_reports: int) -> Dict:
    """虚拟信道比对工作进程函数
    
    这个函数在独立进程中运行，用于并行比对
    """
    try:
        # 在工作进程中创建新的Cython引擎
        engine = get_cython_engine(max_diff_reports, False)  # 工作进程中不使用并行
        
        # 创建临时比对器
        comparer = CythonEnhancedComparer(max_diff_reports, False)
        comparer.cython_engine = engine
        
        # 比对单个虚拟信道
        result = comparer._compare_single_vc_enhanced(vc_id, frames1, frames2, compare_range)
        
        return result
        
    except Exception as e:
        return {
            'vc_id': vc_id,
            'error': str(e),
            'frames_in_base': len(frames1),
            'frames_in_target': len(frames2)
        }


class CythonStreamingComparer:
    """Cython流式比对器
    
    专门用于处理超大文件的流式比对
    """
    
    def __init__(self, chunk_size: int = 5000, max_diff_reports: int = 100):
        """初始化流式比对器
        
        Args:
            chunk_size: 数据块大小
            max_diff_reports: 最大差异报告数量
        """
        self.chunk_size = chunk_size
        self.max_diff_reports = max_diff_reports
        self.logger = logging.getLogger('cython_streaming_comparer')
        
        # 获取Cython引擎
        self.cython_engine = get_cython_engine(max_diff_reports, False)  # 流式处理中不使用进程并行
        
    def compare_large_files_streaming(self,
                                    frames_generator1,
                                    frames_generator2,
                                    total_frames1: int,
                                    total_frames2: int,
                                    compare_range: Optional[Tuple[int, int]] = None,
                                    progress_callback: Optional[Callable] = None) -> Dict:
        """流式比对大文件
        
        Args:
            frames_generator1: 基准文件的帧生成器
            frames_generator2: 目标文件的帧生成器
            total_frames1: 基准文件总帧数
            total_frames2: 目标文件总帧数
            compare_range: 比对范围
            progress_callback: 进度回调
            
        Returns:
            dict: 比对结果
        """
        start_time = time.time()
        
        comparison_result = {
            'total_frames': [total_frames1, total_frames2],
            'frame_results': [],
            'matching_frames': 0,
            'virtual_channel_results': {},
            'processing_time': 0,
            'performance_info': self.cython_engine.get_performance_info()
        }
        
        vc_stats = {}
        processed_frames = 0
        chunk_count = 0
        
        try:
            frames_buffer1 = []
            frames_buffer2 = []
            
            # 流式处理数据块
            for frame1, frame2 in zip(frames_generator1, frames_generator2):
                if frame1:
                    frames_buffer1.append(frame1)
                if frame2:
                    frames_buffer2.append(frame2)
                
                # 当缓冲区达到块大小时，进行比对
                if len(frames_buffer1) >= self.chunk_size or len(frames_buffer2) >= self.chunk_size:
                    # 使用Cython增强比对器处理当前块
                    chunk_comparer = CythonEnhancedComparer(self.max_diff_reports, False)
                    
                    chunk_result = chunk_comparer.compare_files_enhanced(
                        frames_buffer1, frames_buffer2, compare_range, None
                    )
                    
                    # 合并结果
                    self._merge_chunk_results(comparison_result, chunk_result, vc_stats)
                    
                    processed_frames += len(frames_buffer1) + len(frames_buffer2)
                    
                    # 清空缓冲区
                    frames_buffer1.clear()
                    frames_buffer2.clear()
                    chunk_count += 1
                    
                    if progress_callback:
                        progress = int(processed_frames / (total_frames1 + total_frames2) * 100)
                        progress_callback(progress, f"处理数据块{chunk_count}", 
                                        f"已处理{processed_frames}帧")
                    
                    # 内存清理
                    gc.collect()
            
            # 处理剩余的帧
            if frames_buffer1 or frames_buffer2:
                chunk_comparer = CythonEnhancedComparer(self.max_diff_reports, False)
                chunk_result = chunk_comparer.compare_files_enhanced(
                    frames_buffer1, frames_buffer2, compare_range, None
                )
                self._merge_chunk_results(comparison_result, chunk_result, vc_stats)
            
            # 更新最终结果
            comparison_result['virtual_channel_results'] = vc_stats
            comparison_result['processing_time'] = time.time() - start_time
            
            self.logger.info(f"Cython流式比对完成: 处理{processed_frames}帧, 耗时{comparison_result['processing_time']:.2f}秒")
            
            return comparison_result
            
        except Exception as e:
            self.logger.error(f"Cython流式比对失败: {e}", exc_info=True)
            raise
    
    def _merge_chunk_results(self, main_result: Dict, chunk_result: Dict, vc_stats: Dict):
        """合并数据块结果"""
        # 合并匹配帧数
        main_result['matching_frames'] += chunk_result.get('matching_frames', 0)
        
        # 合并帧结果
        main_result['frame_results'].extend(chunk_result.get('frame_results', []))
        
        # 合并虚拟信道结果
        chunk_vc_results = chunk_result.get('virtual_channel_results', {})
        for vc_id, vc_result in chunk_vc_results.items():
            if vc_id not in vc_stats:
                vc_stats[vc_id] = {
                    'vc_id': vc_id,
                    'frames_in_base': 0,
                    'frames_in_target': 0,
                    'matching_frames': 0,
                    'mismatching_frames': 0,
                    'missing_in_target': 0,
                    'missing_in_base': 0,
                    'frame_comparisons': []
                }
            
            # 累加统计信息
            vc_stats[vc_id]['frames_in_base'] += vc_result.get('frames_in_base', 0)
            vc_stats[vc_id]['frames_in_target'] += vc_result.get('frames_in_target', 0)
            vc_stats[vc_id]['matching_frames'] += vc_result.get('matching_frames', 0)
            vc_stats[vc_id]['mismatching_frames'] += vc_result.get('mismatching_frames', 0)
            vc_stats[vc_id]['missing_in_target'] += vc_result.get('missing_in_target', 0)
            vc_stats[vc_id]['missing_in_base'] += vc_result.get('missing_in_base', 0)
            
            # 合并帧比对结果
            vc_stats[vc_id]['frame_comparisons'].extend(vc_result.get('frame_comparisons', [])) 